package org.easy.liblibai;


import org.easy.core.dto.ImgSpec;
import org.easy.core.exception.BizRuntimeException;
import org.easy.core.util.Lists;
import org.easy.liblibai.dto.LibLibGenerateParams;
import org.easy.liblibai.dto.LibLibImg2ImgParams;
import org.easy.liblibai.dto.LibLibTxt2ImgParams;
import org.easy.liblibai.dto.common.LibLibGenerateStatusResponse;
import org.easy.liblibai.dto.enums.LibLibGenerateStatus;

import java.util.List;
import java.util.stream.Collectors;

/**
 * 通用方法封装工具类
 */
public class LibLibAiTemplate {

    private LibLibAiClient libLibAiClient;

    public LibLibAiTemplate(LibLibAiClient client) {
        this.libLibAiClient = client;
    }

    public List<String> genUltraTextToImg(String englishPrompt, ImgSpec imgSpec, Integer imgCount) {
        LibLibGenerateParams libLibGenerateParams = new LibLibGenerateParams();
        libLibGenerateParams.setPrompt(englishPrompt);
        libLibGenerateParams.setImageSize(imgSpec);
        libLibGenerateParams.setImgCount(imgCount != null && imgCount > 0 ? imgCount : 1);
        String generateUuid = libLibAiClient.text2imgUltra(libLibGenerateParams);
        return getAllResult(generateUuid);
    }

    public List<String> genUltraImgToImg(String englishPrompt, String sourceImage, Integer imgCount) {
        LibLibGenerateParams libLibGenerateParams = new LibLibGenerateParams();
        libLibGenerateParams.setPrompt(englishPrompt);
        libLibGenerateParams.setSourceImage(sourceImage);
        libLibGenerateParams.setImgCount(imgCount != null && imgCount > 0 ? imgCount : 1);
        String generateUuid = libLibAiClient.img2imgUltra(libLibGenerateParams);
        return getAllResult(generateUuid);
    }

    public List<String> text2imgAdvance(LibLibTxt2ImgParams libLibImg2ImgParams) {
        String generateUuid = libLibAiClient.text2imgAdvance(libLibImg2ImgParams);
        return getAllResult(generateUuid);
    }

    public List<String> img2imgAdvance(LibLibImg2ImgParams libLibImg2ImgParams) {
        String generateUuid = libLibAiClient.img2imgAdvance(libLibImg2ImgParams);
        return getAllResult(generateUuid);
    }


    private List<String> getAllResult(String generateUuid) {
        LibLibGenerateStatusResponse status = libLibAiClient.status(generateUuid);

        List<Integer> stopLoopCondition = Lists.newArrayList(LibLibGenerateStatus.SUCCESS.getCode(), LibLibGenerateStatus.FAILURE.getCode());
        while (!stopLoopCondition.contains(status.getGenerateStatus())) {
            if (status.getGenerateStatus() == LibLibGenerateStatus.SUCCESS.getCode()) {
                return status.getImages().stream().map(LibLibGenerateStatusResponse.LibLibImageDto::getImageUrl).collect(Collectors.toList());
            } else if (status.getGenerateStatus() == LibLibGenerateStatus.FAILURE.getCode()) {
                throw new BizRuntimeException(status.getGenerateMsg());
            }

            try {
                Thread.sleep(1000);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }

            status = libLibAiClient.status(generateUuid);
        }

        return status.getImages().stream().map(LibLibGenerateStatusResponse.LibLibImageDto::getImageUrl).collect(Collectors.toList());
    }


}